from model import EmbeddedRequestBody, successMessage
from fastapi import APIRouter
from sentence_transformers import SentenceTransformer
from logger import logger

router = APIRouter()
# embeddingModel = SentenceTransformer('C:\\Users\\Administrator\\Desktop\\python\\study\\bigmodel\\moka-ai\\m3e-base')
embeddingModel = SentenceTransformer('bigmodel/moka-ai/m3e-base')


@router.post("/",
             summary="内嵌大模型",
             description="将字符串转化为向量,模型选择为: moka-ai/m3e-base")
def embedding(needEmbeddedStrs: EmbeddedRequestBody):
    try:
        embeddings = embeddingModel.encode(needEmbeddedStrs.strs)
    except Exception as e:
        logger.error(e)
    return successMessage(embeddings.tolist(), "转换成功")
